"""
    TransitionDown Module
"""
import sys, os
sys.path.append(os.path.abspath(os.path.join(os.getcwd())))

import mindspore.nn as nn
import mindspore.ops as ops
import mindspore.numpy as mnp
from mind3d.utils.pointnet2_util import sample_and_group_all, farthest_point_sample, square_distance, index_points


def Square_distance(src, dst):
    out = (src[:, :, None] - dst[:, None]) ** 2
    sum_ = ops.ReduceSum()
    return sum_(out, -1)


def sample_and_group(npoint, radius, nsample, xyz, points, knn=True):
    b, N, c = xyz.shape
    s = npoint
    fps_idx = farthest_point_sample(xyz, npoint)
    new_xyz = index_points(xyz, fps_idx)
    sort = ops.Sort()

    dist = Square_distance(new_xyz, xyz)
    idx = sort(dist)[1][:, :, :nsample]
    grouped_xyz = index_points(xyz, idx)
    grouped_xyz_norm = grouped_xyz - new_xyz.view(b, s, 1, c)

    if points is not None:
        grouped_points = index_points(points, idx)
        new_points = ops.Concat(-1)((grouped_xyz_norm, grouped_points))
    else:
        new_points = grouped_xyz_norm
    return new_xyz, new_points


class PointNet2SetAbstraction(nn.Cell):
    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all, knn):
        super(PointNet2SetAbstraction, self).__init__()
        self.npoint = npoint
        self.radius = radius
        self.nsample = nsample
        self.group_all = group_all
        self.knn = knn
        self.mlp_convs = nn.CellList()
        self.mlp_bns = nn.CellList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1, has_bias=True, weight_init='Uniform', 
                        bias_init = "Uniform"))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel, momentum=0.1, use_batch_statistics = None))
            last_channel = out_channel

        self.relu = ops.ReLU()
        self.transpose = ops.Transpose()
        self.reduce_max = ops.ArgMaxWithValue(axis=2, keep_dims=False)

    def construct(self, xyz, points):
        """SA construct"""

        if self.group_all:
            new_xyz, new_points = sample_and_group_all(xyz, points)
        else:
            new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points, knn=self.knn)

        new_points = self.transpose(new_points, (0, 3, 2, 1))
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points = self.relu(bn(conv(new_points)))
        new_points = self.reduce_max(new_points)[1]
        new_points = self.transpose(new_points, (0, 2, 1))

        return new_xyz, new_points


class Transitiondown(nn.Cell):
    """
    TransitionDown Block
    A key function of TransitionDown Block is to reduce the number of points from the point set as needed.
    Input:
        xyz(Tensor) : Input point sosition data. The shapes of xyz is [B, C, N].
        points(Tensor) : Input position data.The shapes of points is [B, D’, S].
    Output:
        new_xyz: Sampled points position data. The shapes of New_xyz is [B,C,S].
        new_points: Sampled points position feature data. The shapes of New_points is [B,D’,S].
    Argument:
        k: KNN argument.
        n neighbor: KNN argument.
        channels: PointNet++ Set Abstraction module argument.
    Example:
        C = TransitionDown(16,512,32)
    """

    def __init__(self, k, nneighbor, channels):
        super(Transitiondown, self).__init__()
        self.sa = PointNet2SetAbstraction(k, 0, nneighbor, channels[0], channels[1:], group_all=False, knn=True)

    def construct(self, xyz, points):
        """
        TransitionDown construct
        """
        after_sa = self.sa(xyz, points)
        return after_sa